#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <windows.h>

#include "dr_api.h"
#include "util.h"

#define BUF_SIZE 512
#define ENQ 0x05

static char *buf;
static long size;

static file_t log_file;
static char log_buf[BUF_SIZE];

static bool event_pre_syscall(void *, int);
static void event_post_syscall(void *, int);
static void event_exit(void);

extern TCHAR pszFilename[MAX_PATH+1];

DR_EXPORT void dr_init(client_id_t id)
{
	log_file = dr_open_file("C:/syscall.log", DR_FILE_WRITE_OVERWRITE);

	if (log_file == INVALID_FILE)
	{
		dr_messagebox("Unable to create a log file at C:\\");
		exit(1);
	}

    dr_register_pre_syscall_event(event_pre_syscall);
    dr_register_post_syscall_event(event_post_syscall);
    dr_register_exit_event(event_exit);
}

/* Only for Windows XP */
static bool event_pre_syscall(void *drcontext, int sysnum){
	switch (sysnum)
	{
		/* NtCreateFile */
	case 0x25:
		{
			/*
			* Struct/Function prototype we needed.
			*
			* NtCreateFile(
			*	OUT	PHANDLE				FileHandle,
			*	IN	ACCESS_MASK			DesiredAccess,
			*	IN	POBJECT_ATTRIBUTES	ObjectAttributes,		(*)
			*	OUT PIO_STATUS_BLOCK	IoStatusBlock,
			*	IN	PLARGE_INTEGER		AllocationSize	OPTIONAL,
			*	IN	ULONG				FileAttributes,
			*	IN	ULONG				ShareAccess,
			*	IN	ULONG				CreateDisposition,
			*	IN	ULONG				CreateOptions,
			*	IN	PVOID				EaBuffer		OPTIONAL,
			*	IN	ULONG				EaLength
			*	);
			* 
			* struct _OBJECT_ATTRIBUTES{
			*   ULONG					Length;
			*   HANDLE					RootDirectory;
			*	PUNICODE_STRING			ObjectName;				(*)
			*   ULONG					Attributes;
			*   PVOID					SecurityDescriptor;
			*	PVOID					SecurityQualityOfService;
			* } OBJECT_ATTRIBUTES, *POBJECT_ATTRIBUTES;
			* 
			* struct _LSA_UNICODE_STRING {
			*   USHORT					Length;
			*   USHORT					MaximumLength;
			*   PWSTR					Buffer;					(*)
			* } LSA_UNICODE_STRING, *PLSA_UNICODE_STRING, UNICODE_STRING, *PUNICODE_STRING;
			*/
			long *ptemp;

			/* Get POBJECT_ATTRIBUTES */
			ptemp = (long *) dr_syscall_get_param(drcontext, 2);

			/* Get PUNICODE_STRING */
			ptemp = (long *) (*(ptemp+2));

			/* Get PWSTR */
			ptemp = (long *) (*(ptemp+1));

			memset(log_buf, 0, BUF_SIZE);
			sprintf(log_buf, "Open system call, Filename: %ls", (WCHAR *)ptemp);
			size = strlen(log_buf);

			if (!strncmp(log_buf+size - 4, ".txt", 4))
			{
				dr_fprintf(log_file, "%s\n", log_buf);
			}
			
			break;
		}
		/* NtReadFile */
	case 0xb7:
		{
			/*
			* Struct/Function prototype we needed.
			*
			* NtReadFile(
			*   IN	HANDLE				FileHandle,
			*   IN	HANDLE				Event			OPTIONAL,
			*   IN	PIO_APC_ROUTINE		ApcRoutine		OPTIONAL,
			*   IN	PVOID ApcContext	OPTIONAL,
			*   OUT PIO_STATUS_BLOCK	IoStatusBlock,
			*   OUT PVOID				Buffer,					(*)
			*   IN	ULONG				Length,
			*   IN	PLARGE_INTEGER		ByteOffset		OPTIONAL,
			*   IN	PULONG				Key				OPTIONAL
			* );
			*/

			/* Do some preparation for event_post_syscall  */
			GetFileNameFromHandle((HANDLE) dr_syscall_get_param(drcontext, 0));
			buf = (char *) dr_syscall_get_param(drcontext, 5);
			break;
		}
		/* NtWriteFile */
	case 0x112:
		{
			/*
			* Struct/Function prototype we needed.
			*
			* NtWriteFile(
			*   IN	HANDLE				FileHandle,
			*   IN	HANDLE				Event OPTIONAL,
			*   IN	PIO_APC_ROUTINE		ApcRoutine		OPTIONAL,
			*   IN	PVOID				ApcContext		OPTIONAL,
			*   OUT	PIO_STATUS_BLOCK	IoStatusBlock,
			*   IN	PVOID				Buffer,					(*)
			*   IN	ULONG				Length,					(*)
			*   IN	PLARGE_INTEGER		ByteOffset		OPTIONAL,
			*   IN	PULONG				Key				OPTIONAL
			* );
			*/
			/* Content length */
			long len;

			/* Get content */
			buf = (char *) dr_syscall_get_param(drcontext, 5);
			/* Get content length */
			len = (long) dr_syscall_get_param(drcontext, 6);
		
			if (*buf != ENQ)
			{
				/* Write file content */
				dr_fprintf(log_file, "Write system call, Length: %d, Content: %s\n", len, buf);
			}
			else{
				/* ENQ */
				//TODO handle enquiry
			}

			break;
		}
	default:
			break;
	}

	return true;
}

static void event_post_syscall(void *drcontext, int sysnum){
	switch (sysnum)
	{
		/* NtReadFile */
	case 0xb7:
		{
			memset(log_buf, 0, BUF_SIZE);
			sprintf(log_buf, "%ls", pszFilename);
			size = strlen(log_buf);

			if (!strncmp(log_buf+size - 4, ".txt", 4))
			{
				dr_fprintf(log_file, "Read system call, Content: %s\n", buf);
			}

			break;
		}
	default:
			break;
	}
}

static void event_exit(void){
	dr_flush_file(log_file);
	dr_close_file(log_file);
}
